/*
 * Copyright (c) 2011, A. Kandiah
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * 3. Neither the name of the author nor the names of its contributors may be 
 *    used to endorse or promote products derived from this software without 
 *	  specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES,
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY
 * AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL
 * THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
 * OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
 * WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
 * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *
 */

/*
 *  This is the main file of this project. It contains the entrypoint for
 *  the application, as well as the move/rename, undo and reboot functionalities. 
 */


#include <windows.h>
#include <stdio.h>
#include <Strsafe.h>
#include <Winuser.h>

#include <tre/config.h>
#include <tre/regex.h>


#define FROM		0
#define TO			1
#define ERR_MSG_LEN	65536

#define NUM_MATCH	1

#define SESSION_MANAGER_SUBKEY	L"SYSTEM\\CurrentControlSet\\Control\\Session Manager"
#define PENDING_RENAME_VALUE	L"PendingFileRenameOperations"


static void print_error_from_code(DWORD errorCode)
{
    wchar_t message[ERR_MSG_LEN];

    FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, errorCode, 
        MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), message, ERR_MSG_LEN, NULL);

    fwprintf_s(stderr, L"\n");
    fwprintf_s(stderr, L"%s", message);
    fwprintf_s(stderr, L"Error code: %x\n", errorCode);
    exit(errorCode);
}


static void usage(void)
{
    fwprintf_s(stderr, L"\n");
    fwprintf_s(stderr, L"Usage: [OPTION] NAME [NAME]\n");
    fwprintf_s(stderr, L"Moves or removes a file or folder's contents in-use at boot time");
    fwprintf_s(stderr, L"\n");
    fwprintf_s(stderr, L"Option:\n");
    fwprintf_s(stderr, L"   /r\tRestart machine on finish of command\n");
    fwprintf_s(stderr, L"   /u\tUndo scheduling for specified pattern file/folder name.\n");
    fwprintf_s(stderr, L"     \tIf a file is specified, this will remove all scheduled operations for the\n");
    fwprintf_s(stderr, L"     \tspecified file. If a folder is specified, the scheduled operations\n");
    fwprintf_s(stderr, L"     \tare removed from its contents.\n");
    fwprintf_s(stderr, L"   /p\tPrint out the pending rename operations.\n");

    exit(1);
}


static void undo_remove(WCHAR* pattern)
{

    LONG	result;
    HKEY	key;

    WCHAR*	oldValue		= NULL;
    WCHAR*  newValue		= NULL;
    DWORD	dataLength		= 0;

    size_t	oldValueIndex	= 0;
    size_t  newValueIndex	= 0;

    regex_t		preg;
    regmatch_t	pmatch[NUM_MATCH];

    if ((result = RegOpenKeyEx(HKEY_LOCAL_MACHINE, SESSION_MANAGER_SUBKEY, 0, KEY_READ | KEY_SET_VALUE, &key)) 
        != ERROR_SUCCESS)
        print_error_from_code(result);

    // determine the length of the buffer to hold the retrieved value
    if ((result = RegQueryValueEx(key, PENDING_RENAME_VALUE, NULL, NULL, NULL, &dataLength)))
    {
        fprintf(stderr, "\nNo files scheduled for move or deletion\n");
        exit(EXIT_FAILURE);
    }

    oldValue = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY | HEAP_GENERATE_EXCEPTIONS, dataLength);

    if ((result = RegQueryValueEx(key, PENDING_RENAME_VALUE, NULL, NULL, (LPBYTE)oldValue, &dataLength)))
        print_error_from_code(result);

    // allocate the amount of space as the space taken up the existing entries + room for extra null terminator
    newValue = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY | HEAP_GENERATE_EXCEPTIONS, dataLength + sizeof(wchar_t));

    // divide datalength by two inorder to cater for utf-16
    dataLength >>= 1;

    // initialize the regex
    if ((result = regwcomp(&preg, pattern, 0)) != 0) 
    {
        fprintf(stderr, "regcomp() failed, returning nonzero (%d)\n", result);
        exit(EXIT_FAILURE);
    }

    while(oldValueIndex < dataLength - 1)
    {

        // We are going to check if pattern entered by the user matches for the "SOURCE" file/folder. 

        /* NOTE: A REG_MULTI_SZ value is a sequence of zero or more null-terminated strings followed 
        by an additional null-termniator at the end of the sequence.

        For a file-rename operation, the PendingFileRenameOperations registry value will contain 
        the following pattern:	
        OldFileName\0NewFileName\0\0
        For a file-deletion operation, the value for NewFileName will be nothing - represented by 
        only the null terminator. Hence, the value in this instance will be represented as:
        OldFileName\0\0

        Various combinations of the above patterns can occur in the PendingFileRenameOperations 
        registry value. At the end of these patterns, an additional null-terminator must be present.
        */


        // Search for the pattern, expecting one match 
        if (regwexec(&preg, oldValue + oldValueIndex, NUM_MATCH, pmatch, 0) == 0) 
        {
            // Pattern found, so we exclude this string from the updated registry entry. 
            oldValueIndex += wcslen(oldValue + oldValueIndex) +		/* length of the first string */ 
                wcslen(oldValue + oldValueIndex + wcslen(oldValue + oldValueIndex) + 1) + /* length of second string */ 
                1 +													/* space for OldFileName's null terminator */
                1;     										        /* space for NewFileName's null terminator */ 
        } 
        else 
        {

            // Pattern not found, so copy the "source" and "destination"
            // files from the existing registry key to the newly allocated space
            // (which will be written later to the registry)

            CopyMemory(newValue + newValueIndex, 
                oldValue + oldValueIndex, 
                (wcslen(oldValue + oldValueIndex) +					/* length of the first string */ 
                wcslen(oldValue + oldValueIndex + 1 + wcslen(oldValue + oldValueIndex)) + /* length of the second string */ 
                1 +													/* space for OldFileName's null terminator */
                1) << 												/* space for NewFileName's null terminator */ 
                1);													/* multiply by two to cater for UTF-16 */		


            /* The description of the calculation below are as stated above for the CopyMemory call */
            oldValueIndex += wcslen(oldValue + oldValueIndex) + wcslen(oldValue + oldValueIndex + 
                wcslen(oldValue + oldValueIndex) + 1) + 1 + 1;

            newValueIndex += wcslen(newValue + newValueIndex) + wcslen(newValue + newValueIndex + 
                wcslen(newValue + newValueIndex) + 1) + 1 + 1;
        }

    } 

    // If the new length is '0' delete the key
    if (newValueIndex == 0) 
        if ((result = RegDeleteValue(key, PENDING_RENAME_VALUE)))
            print_error_from_code(result);
        else
            goto Cleanup;

    // Write the new value to the registry. NOTE: The addition of '1' to newValueIndex to cater for the extra NULL 
    RegSetValueEx(key, PENDING_RENAME_VALUE, 0, REG_MULTI_SZ, (BYTE*)newValue, (DWORD)(newValueIndex + 1) << 1);  

Cleanup:

    HeapFree(GetProcessHeap(), 0, oldValue);
    HeapFree(GetProcessHeap(), 0, newValue);


    exit(ERROR_SUCCESS);

}


static void print_pending_moves()
{
    LONG	result;
    HKEY	key;

    WCHAR*	data			= NULL;
    DWORD	dataLength		= 0;
    size_t	index			= 0;

    BOOL	source			= TRUE;
    WCHAR	path[MAX_PATH];

    if ((result = RegOpenKeyEx(HKEY_LOCAL_MACHINE, SESSION_MANAGER_SUBKEY, 0, KEY_READ, &key)) != ERROR_SUCCESS)
        print_error_from_code(result);

    // determine the length of the buffer to hold the retrieved value
    if ((result = RegQueryValueEx(key, PENDING_RENAME_VALUE, NULL, NULL, NULL, &dataLength)))
        print_error_from_code(result);

    data = HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY | HEAP_GENERATE_EXCEPTIONS, dataLength + 2);

    if ((result = RegQueryValueEx(key, PENDING_RENAME_VALUE, NULL, NULL, (LPBYTE)data, &dataLength)))
        print_error_from_code(result);

    // divide datalength by two inorder to cater for utf-16
    dataLength >>= 1;

    wprintf(L"\n");
    wprintf(L"Files to be renamed or removed on restart:\n\n");

    // Inorder to deal with unconventional null-termination indicators stored in 
    // the data value of registry entries, we must iterate through the characters 
    // manually. 
    while (index < dataLength - 1)
    {
        if (source)
            wprintf(L"Source:\t");
        else 
            wprintf(L"Target:\t");


        // Check the length of the next string
        if (wcslen(data + index) == 0)
        {
            // The string only had a null-terminator, therefore, 
            // no name was present. Hence the action is to delete the file
            wprintf(L"DELETE\n\n");
            index++;
        }
        else 
        {
            // Retrieve the path from the pattern. More information on the pattern
            // can be found here: http://msdn.microsoft.com/en-us/library/aa365240(v=vs.85).aspx
            swscanf_s((data + index), L"\\??\\%s", path, MAX_PATH);

            // The spacing after the value is different depending on whether
            // the target name or the source name is printed out. 
            (source) ? wprintf(L"%s\n", path) : wprintf(L"%s\n\n", path);

            // Skip the length of the string, the null terminator and 
            // the preceeding question marks and slashes. 
            index += (wcslen(path) + 5);
        }

        source = !source;		
    }


    HeapFree(GetProcessHeap(), 0, data);
}


static void move(LPWSTR from, LPWSTR to)
{

    WCHAR	fromFileName[MAX_PATH];
    HANDLE	findHandle;

    WIN32_FIND_DATA				findData;
    WIN32_FILE_ATTRIBUTE_DATA	fileData;


    // Determine if "from" is a directory
    if (!GetFileAttributesEx(from, GetFileExInfoStandard, &fileData))
        print_error_from_code(GetLastError());

    // If the file is a directory, perform move operation on all the files
    if (fileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
    {
        StringCchCopy(fromFileName, MAX_PATH, from);
        StringCchCat(fromFileName, MAX_PATH, L"\\*");

        if((findHandle = FindFirstFile(fromFileName, &findData)) == INVALID_HANDLE_VALUE)
            print_error_from_code(GetLastError());

        // Go through each file in the directory, and its subdirectory
        do
        {
            // if no files are found
            if (findHandle == INVALID_HANDLE_VALUE)
                return;

            // if a link to current or parent folder is encountered
            if (!wcscmp(findData.cFileName, L".") || !wcscmp(findData.cFileName, L".."))
                continue;

            // create the absolute path to the file - we're going to need it
            StringCbPrintf(fromFileName, MAX_PATH, L"%s\\%s", from, findData.cFileName);

            // if a directory is encountered
            if (findData.dwFileAttributes == FILE_ATTRIBUTE_DIRECTORY)
            {
                move(fromFileName, to);
            }

            // we're here, so it's a file. So, schedule the move operation on the file
            if (!MoveFileEx(fromFileName, to, MOVEFILE_DELAY_UNTIL_REBOOT))
                print_error_from_code(GetLastError());

        } while (FindNextFile(findHandle, &findData));

        FindClose(findHandle);

    } 
    // If a file is given
    else if (!(fileData.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY))
    {
        if (!MoveFileEx(from, to, MOVEFILE_DELAY_UNTIL_REBOOT))
            print_error_from_code(GetLastError());
    }
}


BOOL SetShutDownPrivilege()
{

    HANDLE		currentProcess;
    HANDLE		token;
    LUID		luid;

    TOKEN_PRIVILEGES	tp;


    currentProcess = GetCurrentProcess();

    if (!OpenProcessToken(currentProcess, TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY, &token))
        print_error_from_code(GetLastError());

    // lookup privilege on local system
    if (!LookupPrivilegeValue(NULL, SE_SHUTDOWN_NAME, &luid))        
        print_error_from_code(GetLastError());

    tp.PrivilegeCount = 1;
    tp.Privileges[0].Luid = luid;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    // Enable the privilege or disable all privileges.
    if (!AdjustTokenPrivileges(token, FALSE, &tp, sizeof(TOKEN_PRIVILEGES), (PTOKEN_PRIVILEGES) NULL, 
        (PDWORD) NULL))
        print_error_from_code(GetLastError());

    if (GetLastError() == ERROR_NOT_ALL_ASSIGNED)
        print_error_from_code(GetLastError());

    return TRUE;
}



int wmain(int argc, wchar_t* argv[])
{

    BOOL	restart			= FALSE;
    BOOL	undo			= FALSE;
    DWORD	numFiles		= 0;
    LPWSTR	fileName[2];
    size_t	fileNameLen[2];

    BOOL	isFile			= TRUE;
    LPVOID	fileInformation	= NULL;
    DWORD	errorCode       = ERROR_SUCCESS;

    argc--;
    argv++;

    if (argc == 0 || argc > 3)
        usage();


    SecureZeroMemory(fileNameLen, sizeof(fileNameLen));
    SecureZeroMemory(fileName, sizeof(fileName));

    while(argc > 0)
    {
        if((*(*argv)) == '/' && (*argv)[1] == 'r')
        {
            restart = TRUE;
        } 
        else if ((*(*argv)) == '/' && (*argv)[1] == 'u')
        {
            undo = TRUE;
        }
        else if ((*(*argv)) == '/' && (*argv)[1] == 'p')
        {
            print_pending_moves();
            exit(0);
        }
        else
        {
            if (numFiles > 1)
                usage();

            fileName[numFiles++] = *argv;
        }

        argc--;
        argv++;
    }

    if (numFiles == 0)
        usage();

    if (restart == TRUE && undo == TRUE)
    {
        fprintf(stderr, "Can not undo and restart the machine at the same time\n");
        usage();
    }

    if (undo == TRUE && numFiles == 1)
        undo_remove(fileName[0]);

    // Check if the 'from' path exists 
    if((errorCode = GetFileAttributes(fileName[FROM])) == INVALID_FILE_ATTRIBUTES)
        print_error_from_code(GetLastError());

    // Check that the length of the input path plus 3 is not more than MAX_PATH
    if (fileName[FROM])
        StringCchLength(fileName[FROM], MAX_PATH, &fileNameLen[FROM]);    

    if (fileName[TO])
        StringCchLength(fileName[TO], MAX_PATH, &fileNameLen[TO]);    

    if ((fileNameLen[FROM] > (MAX_PATH - 3)) || (fileNameLen[TO] > (MAX_PATH - 3)))
        print_error_from_code(ERROR_BUFFER_OVERFLOW);

    move(fileName[FROM], fileName[TO]);

    // if we're here, the move has been successful

    // restart if required
    if (restart)
    {
        // adjust the previledges for this process so that it can reboot the system
        SetShutDownPrivilege(); 
        ExitWindowsEx(EWX_REBOOT, SHTDN_REASON_MAJOR_APPLICATION | SHTDN_REASON_MINOR_MAINTENANCE | 
            SHTDN_REASON_FLAG_PLANNED); 

    }

    return ERROR_SUCCESS;
}